Loading Dataset and Annotations

In this notebook we show how to load the data from the StyledCOCO dataset and how to overlay the huma-pose annotations with the images

import os
import sys
import json

import numpy as np
import matplotlib.pyplot as plt
import torchvision

sys.path.append("..")

from data import StyledCoco
from CONFIG import CONFIG

%reload_ext autoreload
%autoreload 2

data_path = CONFIG["paths"]["data_path"]
images_path = os.path.join(data_path, "images", "train")
original_imgs_path = os.path.join(data_path, "original_images", "train2017")
labels_path = os.path.join(data_path, "annotations")
labels_file = os.path.join(labels_path, "person_keypoints_train.json")

dataset = StyledCoco(root=images_path, annFile=labels_file, original_imgs_path=original_imgs_path)

loading annotations into memory...
Done (t=8.47s)
creating index...
index created!

plt.figure(figsize=(11,55))
plt.axis('off')

for i in range(9):
    img, target = dataset.get_styled(i)
    plt.subplot(9,2,2*i+1)
    plt.imshow(img)
    plt.title(f"Styled Image idx {i+1}")
    dataset.coco.showAnns(target)
    
    img, target = dataset.get_original(i)
    plt.subplot(9,2,2*i+2)
    plt.imshow(img)
    plt.title(f"Original Image idx {i+1}")
    dataset.coco.showAnns(target)
    
plt.tight_layout()
plt.show()

This notebook was created by Angel Villar-Corrales